#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : topN_species_rpm.py
@Author  : Your Name
@Date    : 2025-10-15
@Description :
    从 RPM 表中提取丰度最高的前 N 个物种，计算其在样本中的检出率，
    并结合分组信息输出长表格式数据。
"""

import pandas as pd
from argparse import ArgumentParser, Namespace
from pathlib import Path


def main(args: Namespace):
    """
    主函数：读取输入数据、筛选前 N 个物种并生成长表。
    """

    # -----------------------------
    # 1. 创建输出目录
    # -----------------------------
    out_path = Path(args.output_file).resolve()
    out_path.parent.mkdir(parents=True, exist_ok=True)

    # -----------------------------
    # 2. 读取输入文件
    # -----------------------------
    rpm_df = pd.read_table(args.rpm_file, sep="\t")
    group_df = pd.read_table(args.group_file, sep="\t")

    # -----------------------------
    # 3. 计算每个物种的检出样本数
    # -----------------------------
    # 假设第一列为物种名，后续列为样本名
    sample_n = rpm_df.shape[1] - 1
    rpm_df["count"] = (rpm_df.iloc[:, 1:] > 0).sum(axis=1)
    rpm_df["label"] = rpm_df["count"].apply(
        lambda x: f"{x}/{sample_n}({(x / sample_n) * 100:.2f}%)"
    )

    # -----------------------------
    # 4. 按检出样本数排序并取前 N
    # -----------------------------
    rpm_df = rpm_df.sort_values(by="count", ascending=False)
    rpm_df = rpm_df.iloc[:args.top_n, :].drop(columns="count")

    # 记录物种列名（第一列）
    class_name = rpm_df.columns[0]

    # -----------------------------
    # 5. 转换为长表并合并分组信息
    # -----------------------------
    long_df = (
        pd.melt(
            rpm_df,
            id_vars=[class_name, "label"],
            var_name="Sample",
            value_name="RPM",
        )
        .merge(group_df, on="Sample", how="left")
        .query("RPM > 0")  # 仅保留 RPM > 0 的记录
    )

    # -----------------------------
    # 6. 输出结果
    # -----------------------------
    long_df.to_csv(out_path, sep="\t", index=False)
    print(f"输出完成: {out_path}")


if __name__ == "__main__":
    parser = ArgumentParser(description="提取前 N 物种并合并分组信息")
    parser.add_argument("--rpm_file", type=str, required=True, help="输入 RPM 文件路径")
    parser.add_argument("--group_file", type=str, required=True, help="分组信息文件路径")
    parser.add_argument("--output_file", type=str, required=True, help="输出文件路径")
    parser.add_argument("--top_n", type=int, default=10, help="输出前 N 个物种 (默认 10)")
    args = parser.parse_args()
    main(args)
